// Copyright (c) 2011 Bo Zhou<Bo.Schwarzstein@gmail.com>
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//		http://www.apache.org/licenses/LICENSE-2.0 
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/**
 * Refers to "Efficient Acquisition and Realistic Rendering of Car Paint" and "Photo-realistic Rendering of Metallic Car Paint from Image-Based Measurements".
 */

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <shader.h>
#include <float.h>

static const miScalar PI = (miScalar)M_PI;

static inline miScalar clamp(const miScalar& x, const miScalar& min, const miScalar& max)
{
	if (x <= min)
	{
		return min;
	}
	if (x > max)
	{
		return max;
	}
	return x;
}

static inline miScalar geometricalAttenuationFactor(
	const miScalar& dotNH,
	const miScalar& dotNV, 
	const miScalar& dotNL,
	const miScalar& dotVH)
{
	miScalar shadowing = (dotNH * dotNV) / dotVH;
	miScalar masking = (dotNH * dotNL) / dotVH;
	return __min(1.0f, 2.0f * __min(shadowing, masking));
}

static inline miScalar facetSlopeDistribution(
	const miScalar& cosineAlpha,
	const miScalar& roughness)
{
	miScalar clampedCosineAlpha = clamp(cosineAlpha, 0, 1);
	miScalar sinAlpha = sqrtf(1.0f - clampedCosineAlpha * clampedCosineAlpha);
	miScalar tanAlpha = sinAlpha / clampedCosineAlpha;
	miScalar tanAlphaOverRoughness = __max(0, tanAlpha) / roughness;
	return exp(-powf(tanAlphaOverRoughness, 2)) / (powf(roughness, 2) * powf(cosineAlpha, 4));
}

typedef struct {
	miScalar	kD;
	miColor		colorD;

	miScalar	kS0;
	miColor		colorS0;
	miScalar	mS0;
	miScalar	rS0;

	miScalar	kS1;
	miColor		colorS1;
	miScalar	mS1;
	miScalar	rS1;

	miScalar	kS2;
	miColor		colorS2;
	miScalar	mS2;
	miScalar	rS2;

	miScalar	kR;
	miScalar	sampleR;
	miScalar	shininessR;

	miScalar	kFG;

	miInteger	mode;
	int		i_lights;
	int		n_lights;
	miTag		lights[1];
} CarPaint_t;

miC_LINKAGE DLLEXPORT
int CarPaint_version()
{
	return 1;
}

miC_LINKAGE DLLEXPORT
miBoolean CarPaint(
	miColor *result,
	miState *state,
	CarPaint_t *param)
{
	if (state->type == miRAY_SHADOW || state->type == miRAY_DISPLACE)
	{
		return miFALSE;
	}

	miScalar kD = *mi_eval_scalar(&param->kD);
	miColor colorD = *mi_eval_color(&param->colorD);

	miScalar kS0 = *mi_eval_scalar(&param->kS0);
	miColor colorS0 = *mi_eval_color(&param->colorS0);
	miScalar mS0 = *mi_eval_scalar(&param->mS0);
	miScalar rS0 = *mi_eval_scalar(&param->rS0);

	miScalar kS1 = *mi_eval_scalar(&param->kS1);
	miColor colorS1 = *mi_eval_color(&param->colorS1);
	miScalar mS1 = *mi_eval_scalar(&param->mS1);
	miScalar rS1 = *mi_eval_scalar(&param->rS1);

	miScalar kS2 = *mi_eval_scalar(&param->kS2);
	miColor colorS2 = *mi_eval_color(&param->colorS2);
	miScalar mS2 = *mi_eval_scalar(&param->mS2);
	miScalar rS2 = *mi_eval_scalar(&param->rS2);

	miScalar kR = *mi_eval_scalar(&param->kR);
	miInteger sampleR = *mi_eval_integer(&param->sampleR);
	miScalar shininessR = *mi_eval_scalar(&param->shininessR);

	miScalar kFG = *mi_eval_scalar(&param->kFG);

	miInteger mode = *mi_eval_integer(&param->mode);
	miInteger i_lights = *mi_eval_integer(&param->i_lights);
	miInteger n_lights = *mi_eval_integer(&param->n_lights);
	miTag* lights =  mi_eval_tag(param->lights) + i_lights;

	if (mode == 1)
	{
		mi_inclusive_lightlist(&n_lights, &lights, state);
	}
	else if (mode == 2)
	{
		mi_exclusive_lightlist(&n_lights, &lights, state);
	}
	else if (mode == 4)
	{
		n_lights = 0;
		lights = NULL;
	}

	if (mode == 4 || n_lights)
	{
		miVector V = state->dir;
		mi_vector_neg(&V);
		mi_vector_normalize(&V);

		miVector N = state->normal;
		mi_vector_normalize(&N);

		miScalar dotNV = mi_vector_dot(&N, &V);

		for (mi::shader::LightIterator itr(state, lights, n_lights); !itr.at_end(); ++itr)
		{
			miColor resultD;
			resultD.r = resultD.g = resultD.b = resultD.a = 0;

			miColor resultS[3];
			resultS[0].r = resultS[0].g = resultS[0].b = resultS[0].a = 0;
			resultS[1].r = resultS[1].g = resultS[1].b = resultS[1].a = 0;
			resultS[2].r = resultS[2].g = resultS[2].b = resultS[2].a = 0;

			while (itr->sample())
			{
				miScalar dotNL = itr->get_dot_nl();
				if (dotNL <= 0)
				{
					continue;
				}

				miScalar specularNormalizer = 1.0f / (PI * dotNV * dotNL);
				assert(!_isnan(specularNormalizer));

				miVector L = itr->get_direction();
				mi_vector_normalize(&L);
				
				miVector H;
				mi_vector_add(&H, &V, &L);
				mi_vector_normalize(&H);

				miScalar dotVH = mi_vector_dot(&V, &H);
				assert(dotVH > -FLT_EPSILON);

				miScalar dotNH = mi_vector_dot(&N, &H);
				assert(dotVH > -FLT_EPSILON);

				miColor Cl;
				itr->get_contribution(&Cl);

				miScalar G = geometricalAttenuationFactor(dotNH, dotNV, dotNL, dotVH); 
				assert(G > 0);

				if (kD > FLT_EPSILON)
				{
					resultD.r += Cl.r * colorD.r * dotNL;
					resultD.g += Cl.g * colorD.g * dotNL;
					resultD.b += Cl.b * colorD.b * dotNL;
				}

				if (kS0 > FLT_EPSILON)
				{
					miScalar f0 = rS0 + (1 - rS0) * powf(dotVH, 5);
					assert(!_isnan(f0));

					miScalar D0 = facetSlopeDistribution(dotNH, mS0);
					resultS[0].a = f0 * D0 * G * specularNormalizer;
					assert(!_isnan(resultS[0].a));

					resultS[0].r += Cl.r * colorS0.r * resultS[0].a;
					resultS[0].g += Cl.g * colorS0.g * resultS[0].a;
					resultS[0].b += Cl.b * colorS0.b * resultS[0].a;
				}

				if (kS1 > FLT_EPSILON)
				{
					miScalar f1 = rS1 + (1 - rS1) * powf(dotVH, 5);

					miScalar D1 = facetSlopeDistribution(dotNH, mS1);
					resultS[1].a = f1 * D1 * G * specularNormalizer;
					assert(!_isnan(resultS[1].a));

					resultS[1].r += Cl.r * colorS1.r * resultS[1].a;
					resultS[1].g += Cl.g * colorS1.g * resultS[1].a;
					resultS[1].b += Cl.b * colorS1.b * resultS[1].a;
				}

				if (kS2 > FLT_EPSILON)
				{
					miScalar f2 = rS2 + (1 - rS2) * powf(dotVH, 5);

					miScalar D2 = facetSlopeDistribution(dotNH, mS2);
					resultS[2].a = f2 * D2 * G * specularNormalizer;
					assert(!_isnan(resultS[2].a));

					resultS[2].r += Cl.r * colorS2.r * resultS[2].a;
					resultS[2].g += Cl.g * colorS2.g * resultS[2].a;
					resultS[2].b += Cl.b * colorS2.b * resultS[2].a;
				}
			}

			miColor finalResultD;
			finalResultD.r = kD * resultD.r;
			finalResultD.g = kD * resultD.g;
			finalResultD.b = kD * resultD.b;

			miColor finalResultS;
			finalResultS.r = kS0 * resultS[0].r + kS1 * resultS[1].r + kS2 * resultS[2].r;
			finalResultS.g = kS0 * resultS[0].g + kS1 * resultS[1].g + kS2 * resultS[2].g;
			finalResultS.b = kS0 * resultS[0].b + kS1 * resultS[1].b + kS2 * resultS[2].b;

			miInteger samples = itr->get_number_of_samples();
			if (samples)
			{
				miScalar oneDivideSamples = 1.0f / (miScalar)samples;

				finalResultD.r *= oneDivideSamples;
				finalResultD.g *= oneDivideSamples;
				finalResultD.b *= oneDivideSamples;

				finalResultS.r *= oneDivideSamples;
				finalResultS.g *= oneDivideSamples;
				finalResultS.b *= oneDivideSamples;
			}

			result->r += finalResultD.r + finalResultS.r;
			result->g += finalResultD.g + finalResultS.g;
			result->b += finalResultD.b + finalResultS.b;

			assert(!_isnan(result->r));
			assert(!_isnan(result->g));
			assert(!_isnan(result->b));
		}

		if (kR > FLT_EPSILON)
		{
			miColor finalResultR;
			finalResultR.r = finalResultR.g = finalResultR.b = finalResultR.a = 0;

			miScalar rR = (rS0 + rS1 + rS2) / 3;

			double sample[2] = {0, 0};
			int sampleNumber = 0;
			while (mi_sample(sample, &sampleNumber, state, 2, (miUint*)&sampleR))
			{
				miVector R;
				mi_reflection_dir_glossy_x(&R, state, shininessR, sample);

				miVector H;
				mi_vector_add(&H, &V, &R);
				mi_vector_normalize(&H);

				miScalar dotVH = mi_vector_dot(&V, &H);

				miScalar fR = rR + (1 - rR) * powf(dotVH, 5);

				miColor resultR;
				if (mi_trace_reflection(&resultR, state, &R))
				{
					finalResultR.r += resultR.r * fR;
					finalResultR.g += resultR.g * fR;
					finalResultR.b += resultR.b * fR;
				}
			}

			miScalar oneDivideSampleR = 1.0f / sampleR;
			result->r += kR * finalResultR.r * oneDivideSampleR;
			result->g += kR * finalResultR.g * oneDivideSampleR;
			result->b += kR * finalResultR.b * oneDivideSampleR;

			assert(!_isnan(result->r));
			assert(!_isnan(result->g));
			assert(!_isnan(result->b));
		}

		if (kFG > FLT_EPSILON)
		{
			miColor resultFG;
			if (mi_compute_avg_radiance(&resultFG, state, 'f', NULL))
			{
				result->r += kFG * colorD.r * resultFG.r;
				result->g += kFG * colorD.g * resultFG.g;
				result->b += kFG * colorD.b * resultFG.b;
			}

			assert(!_isnan(result->r));
			assert(!_isnan(result->g));
			assert(!_isnan(result->b));
		}
	}


	result->a = 1.0f;

	return miTRUE;
}
